from PIL import Image
from lang_sam import LangSAM

model = LangSAM()


def mask_generate(image_path, prompt, progress):
    """
    生成mask
    """

    print("##########   开始生成mask   ##########")
    # 需要编辑的图片
    image_pil = Image.open(image_path).convert("RGB")

    # 大模型生成
    text_prompt = prompt

    # 结果集
    """Predicts masks for given images and text prompts using GDINO and SAM models.

        Parameters:
            images_pil (list[Image.Image]): List of input images.
            texts_prompt (list[str]): List of text prompts corresponding to the images.
            box_threshold (float): Threshold for box predictions.
            text_threshold (float): Threshold for text predictions.

        Returns:
            list[dict]: List of results containing masks and other outputs for each image.
            Output format:
            [{
                "boxes": np.ndarray,
                "scores": np.ndarray,
                "masks": np.ndarray,
                "mask_scores": np.ndarray,
            }, ...]
    """
    # 生成mask，原理是GINO（目标检测）+SAM（语义分割）
    masks = model.predict([image_pil], [text_prompt])[0]['masks']

    # 保存mask
    for i, mask in enumerate(masks):
        # 将mask转换为PIL图像（二值图）
        mask_image = Image.fromarray(mask.astype(np.uint8) * 255).convert("L")

        # 保存mask
        mask_image.save(f"./masks/mask_{progress}_{i}.png")

        # 将mask叠加到原图上保存（可视化）
        image_with_mask = image_pil.copy()
        image_with_mask.putalpha(mask_image)
        image_with_mask.save(f"./masks/overlay_{progress}_{i}.png")
    print("##########   mask生成结束   ##########")

    # 返回保存的第一张mask图片路径
    return f"./masks/mask_{progress}_0.png"
